function [psi,phi] = cmfd(psi,phi,Q,in,mu,w,dx,g)
% function [psi,phi] = cmfd(psi,phi,Q,in,mu,w,dx)
% This function performs one iteration of coarse mesh finite difference,
%   following the formulation given by Zhong et al. in NSE(158).
%
%   Inputs:
%           psi     -- fine mesh angular flux
%           phi     -- fine mesh scalar flux
%           Q       -- fine mesh external source
%           in      -- inpute structure
%           mtt     -- material index for each fine mesh
%           mu      -- angle set
%           w       -- angle weight
%           dx      -- fine mesh delta x
%   Output
%           psi     -- updated fine mesh angular flux
%           phi     -- updated fine mesh scalar flux

% allocate
N  = length(in.xfm);            % number of coarse meshes
sT = zeros( N, 1 );             % homogenized total cross-section
sS = zeros( N, 1 );             % homogenized scattering cross-section
D  = zeros( N, 1 );             % homogenized diffusion coefficient
Dt = zeros( N+1, 1 );           % "D tilde"
Dh = zeros( N+1, 1 );           % "D hat" (current correction factor)
S  = zeros( N, 1 );             % region source
J  = zeros( N+1, 1 );           % net current at edges
phi_ref  = zeros( N, 1 );       % reference coarse mesh scalar flux

% coarse mesh delta x's
h = in.xcm(2:end)-in.xcm(1:end-1);

% the angle-integrated external source
s = zeros(sum(in.xfm),1);  
for i = 1:length(s)
    s(i) = sum( Q(i,:,g)'.*w(:) );
end

J(1) = sum( mu .* w.* psi(1,:,g)' ); % first edge net current
for i = 1:length(in.xfm)  
    % note, the homogenization is meaningless until I implement a way to
    % specify coarse meshes that contain heterogeneities
    idx1        = 1 + sum( in.xfm(1:(i-1)) );          % lower index
    idx2        = sum( in.xfm(1:(i  )) );              % upper index   
    J(i+1)      = sum( mu .* w.* psi(idx2+1,:,g)' );     % net current at edges
    phi_ref(i)  = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g) ) / h(i);
    sS(i)       = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g)*in.data( in.mt(i), 5 ))/(phi_ref(i)*h(i));
    sT(i)       = sum( dx(idx1:idx2)'.*phi(idx1:idx2,g)*in.data( in.mt(i), 1 ))/(phi_ref(i)*h(i));
    D(i)        = 1/(3*sT(i));
    S(i)        = sum( dx(idx1:idx2)'.*s(idx1:idx2,1) ) / h(i);    
end

a = 0.5;
% note, I am not confident about these Dt and Dh definitions.  For example,
% the cited paper gives Dt = -2*a..., i.e. negative of what I have.  
Dt(1)   = 2*a*D(1) / (a*h(1)+2*D(1));
Dt(end) = -2*a*D(end) / (a*h(end)+2*D(end));
Dh(1)   = -(J(1)+Dt(1)*phi_ref(1))/phi_ref(1);
Dh(end) = -(J(end)+Dt(end)*phi_ref(end))/phi_ref(end);
for i = 2:length(Dh)-1
   Dt(i) = 2*D(i-1)*D(i)/(D(i-1)*h(i)+D(i)*h(i-1));
   Dh(i) = - ( J(i) + Dt(i)*(phi_ref(i)-phi_ref(i-1)) ) / (phi_ref(i)+phi_ref(i-1)) ;
end
% L C U (diagonals of matrix A, for A*phi = S)
C = zeros(N,1);
L = zeros(N-1,1);
U = L;
% left edge
C(1) = 1/h(1) * ( Dt(2) - Dh(2) + Dt(1) + Dh(1) ) + sT(1)-sS(1);
U(1) = 1/h(1) * (-Dt(2)-Dh(2));
% right edge
C(end) = 1/h(end) * ( -Dt(end) + Dt(end-1) - Dh(end) + Dh(end-1) ) + sT(end)-sS(end);
L(end) = 1/h(end) * (-Dt(end-1)+Dh(end-1));
% internal
for i = 2:length(C)-1
    C(i) = 1/h(i) * (Dt(i+1)+Dt(i)-Dh(i+1)+Dh(i)) + sT(i)-sS(i);  % C * phi(i)
    U(i) = 1/h(i) * (-Dt(i+1)-Dh(i+1));                           % U * phi(i+1)
    L(i-1) = 1/h(i) * (-Dt(i)+Dh(i));                             % L * phi(i-1)
end
A = diag(C) + diag(L,-1) + diag(U,1);
phi_c = A\S; % updated coarse mesh scalar flux

% update:  
%    within meshes    
%           psi(ii,m) = psi(ii,m)*f(i),    ii E i
%    add coarse mesh edges
%           psi(ii,m) = psi(ii,m)*f(i),    ii E S(i,i')  and mu > 0
%           psi(ii,m) = psi(ii,m)*f(i'),   ii E (Si,i')  and mu < 0
for i = 1:N  
    % within cell indices
    idx1 = 1 + sum( in.xfm(1:(i-1)) ); % lower index
    idx2 = sum( in.xfm(1:(i  )) );     % upper index   
    phi(idx1:idx2,g)   = phi(idx1:idx2,g)  * phi_c(i)/phi_ref(i);
end

% could also do something with boundary surfaces...perhaps relevant for
% reflective conditions.  Skip for now.

end